from typing import Union

from src.netWork.OOP.Connections import Connections as  Connections
from src.netWork.OOP.ConstNode import ConstNode
from src.netWork.OOP.Layer import Layer as Layer
from src.netWork.OOP.Connection import Connection as Connection
from src.netWork.OOP.Node import Node


class Network(object):

    def __init__(self, lays):

        '''
        initializing a fully connected neural network
        :param lays:
        :param layers: super parameter about how many nodes in each layer in fully connected neural network
        '''

        self.connections = Connections()
        self.layers = []
        layer_count = len(lays)
        node_count = 0
        for i in range(layer_count):
            self.layers.append(Layer(layer_index=i, node_count=lays[i]))

        # 剪掉输出层单独处理, 由于是FC neural network，只需要将所有的节点都接入连接就可以了
        for layer in range(layer_count - 1):
            connections = [Connection(upstream_node, downstream_node)
                           # 上一层用来接受上一层的输出 ai
                           for upstream_node in self.layers[layer].nodes
                           # 下一层用来接受下一次的误差 δj
                           for downstream_node in self.layers[layer + 1].nodes[:-1]]

            # 为每个节点添加上下游集群
            for conn in connections:
                self.connections.add_connection(conn)
                conn.downstream_node.append_upstream_connection(conn)
                conn.upstream_node.append_downstream_connection(conn)

    def train(self, labels, data_set, rate, iteration):
        '''
        训练神经网络
        :param labels: ti
        :param data_set:  x
        :param rate: η
        :param iteration:   times of train
        :return:
        '''
        for i in range(iteration):
            for d in range(len(data_set)):
                self.train_one_sample(labels[d], data_set[d], rate)

    def train_one_sample(self, label, sample, rate):
        '''
        inner function to train a network for a sample
        :param labels:
        :param sample:
        :param rate:
        :return:
        '''
        self.predict(sample)
        self.calc_delta(label)
        self.update_weight(rate)

    def calc_delta(self, label):
        '''
        inner function to calculate each node's delta
        :param label:
        :return:
        '''
        output_nodes = self.layers[-1].nodes
        for i in range(len(label)):
            output_nodes[i].calc_output_layer_delta(label[i])
        # 从倒数第二个开始倒序排列，直到最后一层截至 for example: [1, 2, 3, 4, 5] to [4, 3, 2, 1]
        for layer in self.layers[-2::-1]:
            # 相当于从后往前计算隐藏层的 δk
            for node in layer.nodes:
                node.calc_hidden_layer_delta()

    def update_weight(self, rate):
        '''
        inner function to update each node's weight
        :param rate:
        :return:
        '''
        for layer in self.layers[:-1]:
            for node in layer.nodes:
                for conn in node.downstream:
                    conn.update_weight(rate)

    def calc_gradient(self):
        '''
        inner function to calculate each connection's gradient
        :return:
        '''
        for layer in self.layers[:-1]:
            node: Node
            for node in layer.nodes:
                conn: Connection
                for conn in node.downstream:
                    conn.calc_gradient()

    def get_gradient(self, label, sample):
        '''
        获取网络下的一个样本，每个连接上的梯度
        :param label:
        :param sample:
        :return:
        '''
        self.predict(sample)
        self.calc_delta(label)
        self.calc_gradient()

    def predict(self, sample):
        '''
        根据
        :param sample:
        :return:
        '''
        self.layers[0].set_output(sample)
        for i in range(1, len(self.layers)):
            self.layers[i].calc_output()
        return map(lambda node: node.output, self.layers[-1].nodesp[:-1])

    def dump(self):
        '''

        :return:
        '''
        for layer in self.layers:
            layer.dump()